import json
from typing import Dict
from optimize import extractPdfInfo


def load_training_data(data_path, max_samples=None):
    """加载训练数据"""
    data = []
    with open(data_path, 'r', encoding='utf-8') as f:
        for i, line in enumerate(f):
            if max_samples and i >= max_samples:
                break
            try:
                data.append(json.loads(line.strip()))
            except json.JSONDecodeError:
                print(f"跳过无效JSON行: {i + 1}")
                continue

    print(f"📊 加载了 {len(data)} 个训练样本")
    return data


def validate_extraction(extracted: Dict, expected: Dict) -> Dict:
    """验证提取结果的准确性"""
    validation_result = {}

    for field in expected.keys():
        if field in extracted:
            if extracted[field] == expected[field]:
                validation_result[field] = "正确"
            else:
                validation_result[field] = f"错误: 期望 {expected[field]}, 得到 {extracted[field]}"
        else:
            validation_result[field] = "缺失"

    accuracy = sum(1 for v in validation_result.values() if v == "正确") / len(validation_result)
    validation_result["总体准确率"] = f"{accuracy:.2%}"

    return validation_result


# 测试验证
expected_result = {
    "MBL/HBL No.": "149506063292",
    "Container No.": "GAOU6252103",
    "Date of Exportation": "2025/09/18",  # 注意: 实际日期是2025/09/18
    "Shipper/Exporter[Company]": "SHENZHEN ZONGYU TECHNOLOGY CO., LTD.",
    "Shipper/Exporter[Address]": "ROOM 403, RENHE XINGYE BUSINESS CENTER, 24-1 NANLIAN ROAD, LONGGANG DISTRICT, SHENZHEN ZIP CODE:510000",
    "Consignee[Company]": "DONGJOO TRADE INC",
    "Consignee[Address]": "3460 W 7TH ST APT 809 LOS ANGELES, CA 90005",
    "packages": 936,
    "weight": 12369,
    "volume": 68.1
}



if __name__ == '__main__':
    # 加载数据
    print("📊 加载训练数据...")
    train_data = load_training_data(r'D:\存档-可导入/dataset_info.jsonl', max_samples=5)  # 先用5条测试
    for data in train_data:
        result = extractPdfInfo.extractPdf(data['text'])
        # 验证结果
        validation = validate_extraction(result, data['labels'])
        print("\n验证结果:")
        print(json.dumps(validation, indent=2, ensure_ascii=False))